import pandas as pd
import matplotlib.pyplot as plt
import numpy as np

data = pd.read_excel('2021MCMProblemC2.xlsx')
data = data[data['Lab Status']!='Negative ID']

fitData=[]
x=data['Latitude'].tolist()
y=data['Longitude'].tolist()

for i in range(len(x)):
    fitData.append([x[i],y[i]])
fitData=np.array(fitData)


from sklearn.cluster import KMeans

kmeans=KMeans(n_clusters=4)
kmeans.fit(fitData)


def visual_kmeans_effect(k_means, dataset):
    assert dataset.shape[1] == 2, 'only support dataset with 2 features'
    X = dataset[:, 0]
    Y = dataset[:, 1]
    X_min, X_max = np.min(X) - 1, np.max(X) + 1
    Y_min, Y_max = np.min(Y) - 1, np.max(Y) + 1
    # meshgrid 生成网格点坐标矩阵
    X_values, Y_values = np.meshgrid(np.arange(X_min, X_max, 0.01),
                                     np.arange(Y_min, Y_max, 0.01))
    # 预测网格点的标记
    predict_labels = k_means.predict(np.c_[X_values.ravel(), Y_values.ravel()])
    predict_labels = predict_labels.reshape(X_values.shape)
    plt.imshow(predict_labels, interpolation='nearest',
               extent=(X_values.min(), X_values.max(),
                       Y_values.min(), Y_values.max()),
               cmap=plt.cm.Paired,
               aspect='auto',
               origin='lower')
    '''
    # 将数据集绘制到图表中
    plt.scatter(X, Y, marker='v', facecolors='none', edgecolors='k', s=30)

    # 将中心点绘制到图中
    centroids = k_means.cluster_centers_
    plt.scatter(centroids[:, 0], centroids[:, 1], marker='o',
                s=100, linewidths=2, color='k', zorder=5, facecolors='b')
    plt.title('Clustering Results of Uncertain Honeycomb')
    plt.xlim(X_min, X_max)
    plt.ylim(Y_min, Y_max)
    plt.xlabel('Latitude')
    plt.ylabel('Longitude')
    '''
    plt.show()


visual_kmeans_effect(kmeans,fitData)
